from datetime import datetime, timedelta

import pandas as pd

from app_config import get_engine_ts
from app_config import get_pro
import time

"""
i 上证50.py
"""
# # 获取今天的日期
# today = datetime.now()
"""     1 样本空间    """
''' 1.1 上市时间超过 12 个月； '''
# 指定日期为2024年5月31号
end_date = datetime(2024, 4, 30)
# 计算一年前的日期
start_date = end_date - timedelta(days=365)
# 按照指定格式输出
start_date_str = start_date.strftime("%Y%m%d")
# 按照指定格式输出
end_date_str = end_date.strftime("%Y%m%d")
print("一年前的今天是:", start_date_str)

counter = {'count': 0}


def get_profit(st_code):
    counter['count'] += 1
    print(counter['count'])
    if counter['count'] % 190 == 0:
        print(f"Sleeping for 60 seconds after 150 requests...")
        time.sleep(62)

    df_ = get_pro().query('fina_indicator', ts_code=st_code,
                          fields="ts_code,end_date,q_opincome,q_investincome,q_dtprofit,q_eps,q_netprofit_margin,"
                                 "q_gsprofit_margin,q_exp_to_sales,q_profit_to_gr,q_saleexp_to_gr,q_adminexp_to_gr,"
                                 "q_finaexp_to_gr,q_impair_to_gr_ttm,q_gc_to_gr,q_op_to_gr,q_roe,q_dt_roe,q_npta,"
                                 "q_opincome_to_ebt,q_investincome_to_ebt,q_dtprofit_to_profit,q_salescash_to_or,"
                                 "q_ocf_to_sales,q_ocf_to_or", start_date=start_date_str,
                          end_date=end_date_str)

    df_['q_dtprofit_to_profit'] = df_['q_dtprofit_to_profit'].fillna(100)

    df_['q_profit'] = df_['q_dtprofit'] / df_['q_dtprofit_to_profit'] * 100
    print("财务数据 " + st_code + "  :  " + str(len(df_)))
    if len(df_) > 4:
        print("err =================" + str(len(df_)))
    df_dedup = df_.drop_duplicates(subset='end_date', keep='first')
    total_q_dtprofit = df_dedup['q_dtprofit'].sum()
    total_q_profit = df_dedup['q_profit'].sum()
    return total_q_dtprofit / 100000000, total_q_profit / 100000000


def test():
    engine = get_engine_ts()

    ttm_date_str = "20240429"

    pro = get_pro()
    st_filtered = pro.index_weight(index_code='000010.SH', trade_date='20240531')
    # 假设df是你的DataFrame
    st_filtered.rename(columns={'con_code': 'ts_code'}, inplace=True)
    stock_name = pro.stock_basic(exchange='', list_status='L',
                                 fields='ts_code,name,symbol')
    st_filtered = pd.merge(st_filtered, stock_name, on='ts_code', how='left')

    """     2、选样方法    """
    '''2.1对样本空间内的证券按照过去一年的日均成交金额由高到低排名，剔除排名后 10%的证券作为待选样本；'''
    # 转换 ts_code 列为元组，并生成 SQL 中 IN 子句所需的格式
    ts_codes = st_filtered['ts_code'].tolist()

    # 构建查询语句
    query = f"""
    SELECT * FROM `daily_basic`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
    AND trade_date >= '{start_date_str}'
    AND trade_date <= '{end_date_str}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)

    result_df['amount'] = result_df['circ_mv'] * result_df['turnover_rate']

    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', 'avg_amount']
    merge = pd.merge(st_filtered, grouped_df, on='ts_code', how='left')

    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', 'avg_total_mv']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    # # 按 avg_amount 列降序排序
    # df_sorted = pd_merge.sort_values(by='avg_amount', ascending=False)
    # # 获取前 90% 的数据
    # num_rows = len(df_sorted)
    # top_50_percent = df_sorted.head(int(num_rows * 0.9))

    top_50_percent = pd_merge.sort_values('avg_total_mv', ascending=False).reset_index()
    # d = top_50_percent.head(50)
    d = top_50_percent
    # top_50_percent.to_excel("top_300.xlsx")
    '''计算TTM'''
    # 使用 apply 方法添加两列 total_q_dtprofit 和 total_q_profit
    # todo
    d[['total_q_dtprofit', 'total_q_profit']] = d['ts_code'].apply(lambda _ts_code: pd.Series(get_profit(_ts_code)))

    # 构建查询语句
    query_date = f"""
         SELECT ts_code, total_mv FROM `daily_basic`
         WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
         AND trade_date = '{ttm_date_str}'
         """

    # 执行查询并将结果转换为DataFrame
    result_df_date = pd.read_sql_query(query_date, engine)
    result_df_date['total_mv'] = result_df_date['total_mv'] / 10000

    data_frame_d = pd.merge(d, result_df_date, on='ts_code', how='left')

    # 确认列名无冲突
    print(data_frame_d.columns)
    data_frame_d['ttm'] = data_frame_d['total_mv'] / data_frame_d['total_q_profit']

    data_frame_d['dt_ttm'] = data_frame_d['total_mv'] / data_frame_d['total_q_dtprofit']

    data_frame_d.to_excel('上证50.xlsx')


if __name__ == '__main__':
    test()
